#!/usr/bin/env python3
import argparse
import os
import sys
import subprocess
import tempfile
import re
import json
import mysql.connector
from mysql.connector import pooling
from dotenv import load_dotenv
import pymysql

'''参考：2936035784322449408
'''

# 加载系统环境变量
load_dotenv()

# 命令行参数解析
parser = argparse.ArgumentParser(
    description="Claude与MySQL交互的中间层"
)
parser.add_argument("-q", "--query", help="直接执行SQL查询")
parser.add_argument("-a", "--ask", help="向Claude提问并执行返回的SQL查询")
parser.add_argument("-i", "--interactive", action="store_true", help="交互式模式")
args = parser.parse_args()


# 连接数据库
def connect():
    conn = pymysql.connect(
        host=os.environ.get("DB_HOST", "localhost"),
        user='root',
        password='123456',
        port=int(os.environ.get("DB_PORT", "3306")),
        database='test'
    )
    return conn


# 执行SQL查询
def execute_query(query):
    # 建立连接
    conn = connect()
    cursor = conn.cursor(pymysql.cursors.DictCursor)
    try:
        cursor.execute(query)
        results = cursor.fetchall()
        return results
    except Exception as e:
        conn.rollback()
        print(e)
        return {"error": str(e)}
    finally:
        cursor.close()
        conn.close()


# 提取Claude回复中的SQL查询
def extract_sql_from_claude_response(response):
    sql_regex = r"```sql\s*([\s\S]*?)\s*```"
    match = re.search(sql_regex, response)
    return match.group(1).strip() if match else None


# 向Claude提问
def ask_claude(question):
    try:
        with tempfile.NamedTemporaryFile(delete=False, mode="w+", encoding="utf-8") as tf:
            temp_file = tf.name
        # 用 shell 调用 claude 命令
        # subprocess.run(f'claude -p "{question}" > "{temp_file}"', shell=True, check=True)
        subprocess.run([f'claude -p "{question}" > "{temp_file}"'], shell=True, check=True)
        with open(temp_file, "r", encoding="utf-8") as f:
            claude_response = f.read()
        os.remove(temp_file)
        return claude_response
    except Exception as e:
        return f"Error asking Claude: {e}"


# 交互式模式
def interactive_mode():
    print("欢迎使用Claude-MySQL交互工具 (输入'exit'退出)")
    try:
        while True:
            question = input("\n请输入您的问题: ").strip()
            if question.lower() == "exit":
                print("已退出。")
                break

            print("正在向Claude提问...")
            claude_response = ask_claude(question)
            print("\nClaude的回答:")
            print(claude_response)

            sql_query = extract_sql_from_claude_response(claude_response)
            if sql_query:
                print("\n检测到SQL查询，正在执行...")
                print(f"SQL: {sql_query}")
                results = execute_query(sql_query)
                print("\n查询结果:")
                print(json.dumps(results, ensure_ascii=False, indent=2))

                # 将结果再次发给Claude进行解释
                print("\n让Claude解释结果...")
                explanation = ask_claude(
                    f'以下是我之前问题"{question}"的SQL查询结果:\n{json.dumps(results, ensure_ascii=False)}\n请解释这些结果。'
                )
                print("\nClaude的解释:")
                print(explanation)
            else:
                print("Claude的回答中没有检测到SQL查询")
    except KeyboardInterrupt:
        print("\n交互模式已退出。")
        sys.exit(0)


def main():
    try:
        if args.query:
            # 直接SQL查询
            results = execute_query(args.query)
            print(json.dumps(results, ensure_ascii=False, indent=2))
        elif args.ask:
            print("正在向Claude提问...")
            claude_response = ask_claude(args.ask)
            print("Claude的回答:")
            print(claude_response)

            sql_query = extract_sql_from_claude_response(claude_response)
            if sql_query:
                print("\n检测到SQL查询，正在执行...")
                print(f"SQL: {sql_query}")
                results = execute_query(sql_query)
                print("\n查询结果:")
                print(json.dumps(results, ensure_ascii=False, indent=2))
                # Claude解释
                print("\n让Claude解释结果...")
                explanation = ask_claude(
                    f'以下是我之前问题"{args.ask}"的SQL查询结果:\n{json.dumps(results, ensure_ascii=False)}\n请解释这些结果。'
                )
                print("\nClaude的解释:")
                print(explanation)
            else:
                print("Claude的回答中没有检测到SQL查询")
        elif args.interactive:
            interactive_mode()
        else:
            parser.print_help()

    except Exception as e:
        print(f"错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
